#include <istream>
#include <vector>
#include <string>
#include <fstream>
#include <iostream>
#include <algorithm>
#include <map>
#include <pthread.h>
#include <boost/format.hpp>
#include <stdlib.h>

#include "UPGMA_nearest_neighbor.h"
#include "KITSCH_nearest_neighbor.h"
#include "bcpp_tree.h"
#include "portable_timer.h"
#include "perlutil.h"
#include "matrix.h"
#include "get_arg.h"
#define DEBUG_LEVEL 2
#include "debug_func.h"

using boost::format;
using std::vector;
using std::ostringstream;

using std::string;
using std::fill;
using std::ofstream;
using std::cout;
using std::cin;
using std::cerr;
using std::ifstream;
using std::map;
using std::pair;
using std::istream;
using std::max;
using std::min;

//________________________________________________________________________________________

//	scale_data

//________________________________________________________________________________________
void scale_data(t_matrix& matrix, const vector<string>& names,
				double ortholog_ks, double scaling_factor,
				const string& tax_name2)
{
	// constants
	double tax2_ks = (scaling_factor * ortholog_ks) / ( 1.0 + scaling_factor);
	double scaling_factor_inv = 1.0 / scaling_factor;


	// figure out which genes belong to species #2
	vector<double>	is_second_taxa(names.size(), 0.0);
	for (unsigned i = 0; i < names.size(); ++i)
		is_second_taxa.push_back((names[i].find(tax_name2) != string::npos) ? 1.0 : 0.0);

	// scale distance in Ks units from ortholog node to tax2
	for (unsigned i = 0; i < names.size() - 1 ; ++i)
		for (unsigned j = i + 1; j < names.size(); ++j)
		{
			double has_second_taxa_in_pair = is_second_taxa[i] + is_second_taxa[j];

			// scale ks for different orthologs
			double ks = matrix[i][j];
			double scale = std::min(ks, tax2_ks * has_second_taxa_in_pair);
			double delta = ks - scale;
			scale *= scaling_factor_inv;
			matrix[i][j] = matrix[j][i] = scale + delta;
		}
}


//________________________________________________________________________________________

//	read_matrix_from_stream

//________________________________________________________________________________________
void dump_matrix_to_stream(	ostream& outfile,
							t_matrix& matrix,
							vector<string>& names)
{
	outfile << string(15, ' ');
	for (unsigned i = 0; i < matrix.size(); ++i)
		outfile << boost::format ("%|1$-15|") % names[i];
	outfile << "\n";

	for (unsigned i = 0; i < matrix.size(); ++i)
	{
		outfile << boost::format ("%|1$-15|")  % names[i];
		for (unsigned j = 0; j < matrix.size(); ++j)
			outfile << boost::format ("%|1$-15|") %
							matrix[i][j];
		outfile << "\n";
	}
	outfile << "\n";
}


//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	print and analyse orthologs

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
bool open_filename(ofstream& ofs, const string& file_name)
{
	ofs.open(file_name.c_str());
	if (ofs)
		return true;
	cerr << "Error\nCould not open " << file_name << "\n";
	return false;
}
bool print_gene_set(ofstream& ofs, const string& file_name, const string& header, const bcpp_tree::t_ortholog_data& data)
{
	if (!data.ensg_sets.size())
		return true;

	if (!ofs.is_open() &&
		!open_filename(ofs, file_name))
		return false;

//	cerr << boost::format("%|2$+3| sets from %|1$-10|:\n\t%|3$+5| tax1 genes\n\t%|4$+5| tax2 genes\n")
//	% file_name % data.ensg_sets.size() % data.cnt_tax1 % data.cnt_tax2;
	ofs << header;
	for (unsigned i = 0; i < data.ensg_sets.size(); ++i)
	{
		for (unsigned j = 0 ; j < data.ensg_sets[i].size(); ++j)
			ofs << data.ensg_sets[i][j].gene << ":"
				<< data.ensg_sets[i][j].len << "\t";
		ofs << "\n";
	}
	ofs << std::endl;
	return true;
}


struct t_outputstreams
{

	t_outputstreams(const string& prefix) : file_prefix(prefix)
	{
		open();

		ofs_all_nodes	<< "To leaf\tTo root\tGene id\tNode type\n"
						<< string(80, '-') << "\n";

		ofs_full_tree	<< "Branch length\tTo root\tTo leaf\tGene id\tNode type\n"
						<< string(80, '-') << "\n";
	}
	string	 file_prefix;
	ofstream timing;
	ofstream ortho_1_to_1;
	ofstream ortho_1_to_m;
	ofstream ortho_m_to_1;
	ofstream ortho_m_to_m;
	ofstream orphaned_1_to_1;
	ofstream orphaned_1_to_m;
	ofstream orphaned_m_to_1;
	ofstream orphaned_m_to_m;
	ofstream orphaned_tax1;
	ofstream orphaned_tax2;
	ofstream ofs_full_tree;
	ofstream newhampshire_tree;
	ofstream ofs_all_nodes;
	bool open();
	bool printout_analysis(	bcpp_tree::t_tree&		tree,
						const string&				header);
};


bool t_outputstreams::printout_analysis(bcpp_tree::t_tree& tree, const string& header)
{
//	cerr << "\n" << string(80, '*') << "\n\n";

	// dump all nodes
	{
		ofs_all_nodes << header << std::endl;
		bcpp_tree::do_dump_nodes dump_nodes(ofs_all_nodes);
		tree.for_each(dump_nodes);
		ofs_all_nodes << std::endl;
	}

	// print full tree
	{

		ofs_full_tree << header << std::endl;
		bcpp_tree::do_print_tree_full print_tree(ofs_full_tree);
		tree.for_each(print_tree);
		ofs_full_tree << std::endl;


		newhampshire_tree << header << std::endl;
		bcpp_tree::do_newhampshire_format_str nh_str(newhampshire_tree);
		tree.for_each(nh_str);
		newhampshire_tree << std::endl;
	}

	// print gene sets
	{
		bcpp_tree::do_collect_orthologs data;
		tree.for_each(data);
		if (!print_gene_set(ortho_1_to_1	, file_prefix + "ortho_1_to_1.data"     ,
													header, data.ortho_1_to_1		) ||
			!print_gene_set(ortho_1_to_m	, file_prefix + "ortho_1_to_m.data"     ,
													header, data.ortho_1_to_m   	) ||
			!print_gene_set(ortho_m_to_1	, file_prefix + "ortho_m_to_1.data"     ,
													header, data.ortho_m_to_1   	) ||
			!print_gene_set(ortho_m_to_m	, file_prefix + "ortho_m_to_m.data"     ,
													header, data.ortho_m_to_m   	) ||
			!print_gene_set(orphaned_1_to_1	, file_prefix + "orphaned_1_to_1.data"  ,
													header, data.orphaned_1_to_1	) ||
			!print_gene_set(orphaned_1_to_m	, file_prefix + "orphaned_1_to_m.data"  ,
													header, data.orphaned_1_to_m	) ||
			!print_gene_set(orphaned_m_to_1	, file_prefix + "orphaned_m_to_1.data"  ,
													header, data.orphaned_m_to_1	) ||
			!print_gene_set(orphaned_m_to_m	, file_prefix + "orphaned_m_to_m.data"  ,
													header, data.orphaned_m_to_m	) ||
			!print_gene_set(orphaned_tax1	, file_prefix + "orphaned_tax1.data"    ,
													header, data.orphaned_tax1  	) ||
			!print_gene_set(orphaned_tax2	, file_prefix + "orphaned_tax2.data"    ,
													header, data.orphaned_tax2  	))
		return false;
	}
	return true;

}



bool t_outputstreams::open()
{
	return(
			open_filename(ofs_full_tree, 		file_prefix + "full_topology.tree"	)	&&
			open_filename(newhampshire_tree,	file_prefix + "new_hampshire.tree"	)	&&
			open_filename(ofs_all_nodes, 		file_prefix + "all_nodes.list"		)	&&
			open_filename(timing, 				file_prefix + "timing"		)

		);
}

struct incorrect_pairs
{
	unsigned cluster_num;
	unsigned expected_pairs;
	unsigned missing_pairs;
	unsigned duplicate_pairs;
	incorrect_pairs(unsigned c, unsigned e, unsigned m, unsigned d):
		cluster_num(c), expected_pairs(e), missing_pairs(m), duplicate_pairs(d){}
};


//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
int main (int argc, char *argv[])
{
	// we use cout to print progress so flush per call
	cout.setf(std::ios_base::unitbuf);

	std::ios::sync_with_stdio(false);

	t_program_args args;
	if (!process_command_line_options(argc, argv, args))
		return 1;


	//
	// open input file
	//
	istream* datafile = &cin;
	ifstream infile;
	if (args.ks_file.length())
	{
		infile.open(args.ks_file.c_str());
		if (!infile)
		{
			std::cerr << "Error opening [" << args.ks_file << "] \n";
			return 1;
		}
		datafile = &infile;
	}

	//
	//	prepare files for output
	//
	t_outputstreams strms (args.out_file_prefix);
	strms.timing <<	"Genes\tSum of squares\tAverage %stddev\tCount Trees\tbest iteration\tTime\n"
					"-----\t--------------\t---------------\t-----------\t--------------\t---------\n";


	unsigned cnt_cluster = 0;
	//
	//	save clusters with missing entries
	//
	vector<incorrect_pairs> problems;



	{
		t_progress_indicator dots(cerr, 5000);
		while (datafile->good())
		{
			++cnt_cluster;


			//
			//	read data from input file
			//
			t_matrix		matrix;
			vector<string> 	names;
			string			cluster_name;
			unsigned		cnt_missing_values;
			unsigned		cnt_duplicate_values;
			t_matrix_bool	missing_values;

			unsigned		cnt_pairs =  read_matrix_from_stream_in_fasta_pairs_format
														(*datafile, matrix,
																	missing_values,
																	names,
																	cluster_name,
																	cnt_missing_values,
																	cnt_duplicate_values);

			if (!cnt_pairs)
				break;

			cluster_name = ">Cluster# " + cluster_name;

			// save missing pairs
			unsigned expected_pairs = names.size() * (names.size() - 1) / 2;
			if (cnt_pairs != expected_pairs)
				problems.push_back(incorrect_pairs(cnt_cluster, expected_pairs,
														cnt_missing_values, cnt_duplicate_values));

			//
			//	scale data if necessary
			//
			if (args.scale_data)
				scale_data(matrix, names, args.speciation, args.clock_adjust, args.taxon2);

			strms.timing << boost::format
										("%|1$-5|\t")
										% names.size();
			strms.timing.flush();


			double sum_sqrs				= 0.0;
			double avg_percent_stddev	= 0.0;
			unsigned cnt_trees_examined = 1;
			unsigned best_iteration		= 0;
			t_portable_timer	timer;
			bcpp_tree::t_tree paralog_tree;
			if (args.analysis_type == eUPGMA || args.analysis_type == eUPGMA_KITSCH)
			{
				upgma_tree::upgma_make_tree(paralog_tree, names, matrix,
											args.fill_lacunae, dots);
				// get sum squares for upgma
				bcpp_tree::do_sum_sqrs analyse_tree(matrix);
				paralog_tree.for_each(analyse_tree);
				analyse_tree.results(sum_sqrs, avg_percent_stddev);
			}
			else if (args.analysis_type == eKITSCH)
				// kitsch make and evaluate
				kitsch_tree::kitsch_make_tree(paralog_tree, names, matrix,
												args.cnt_random_global_searches,
												args.power,
												best_iteration,
												sum_sqrs, avg_percent_stddev,
												cnt_trees_examined,
												dots);
			// optimize UPGMA tree
			if (args.analysis_type == eUPGMA_KITSCH)
			{
				kitsch_tree::optimize_bcpp_tree(paralog_tree, names, matrix,
													args.cnt_random_global_searches,
													args.power,
													args.keyhole_size,
													best_iteration,
													sum_sqrs,
													avg_percent_stddev,
													cnt_trees_examined,
													dots);
			}
			// get sum squares for upgma
			bcpp_tree::do_sum_sqrs2 analyse_tree(matrix);
			paralog_tree.for_each(analyse_tree);
			analyse_tree.results(sum_sqrs, avg_percent_stddev);

			strms.timing << boost::format(	"%|1$14|\t"
											"%|2$15|\t"
											"%|3$11|\t"
											"%|4$14|\t"
											"%|5$9|\n")
											% sum_sqrs
											% avg_percent_stddev
											% cnt_trees_examined
											% best_iteration
											% timer.reset();

			// set type for each node
			bcpp_tree::do_set_node_types set_node_types(args.taxon1, args.taxon2);
			paralog_tree.for_each(set_node_types);

			if (!strms.printout_analysis(paralog_tree, cluster_name))
				return 0;
		}
	}


	//
	//	print clusters with missing entries
	//
	if (problems.size())
	{
		cerr << "Some matrixes had missing or duplicate pairs:\n";
		for (unsigned i = 0; i < problems.size(); ++i)
		{
			cerr	<< format("\tCluster #%|1$+4|: "
								"%|2$+4| pairs missing, "
								"%|3$+4| pairs duplicated "
								"out of %|4$+4| pairs.\n")
					% problems[i].cluster_num
					% problems[i].missing_pairs
					% problems[i].duplicate_pairs
					% problems[i].expected_pairs;
		}
	}
	cerr << "\nAll finished!!\n\n";

	return 0;
}

